#!/usr/bin/env python
#
# Copyright (C) 2011 Clement Lorteau <clem@lorteau.fr>	
#		
# This file is part of QuickScreenShots.		
# 	
# This program is free software; you can redistribute it and/or modify	
# it under the terms of the GNU General Public License as published by		
# the Free Software Foundation; version 3.	
# 	
# This program is distributed in the hope that it will be useful,		
# but WITHOUT ANY WARRANTY; without even the implied warranty of		
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the		
# GNU General Public License for more details.		
# 		
# You should have received a copy of the GNU General Public License		
# along with this program; if not, write to the Free Software		
# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301, USA.

import platform
import os
import wx
from tempfile import NamedTemporaryFile
from PIL import Image

def snap(showcursor = True):
	if platform.system() == 'Windows':
		return snap_win(showcursor)
	elif platform.system() == 'Darwin':
		return snap_osx(showcursor)
	else:
		raise Exception('Platform unrecognized')

def snap_osx(showcursor = True):
	"""
	Hints for multi-screen: http://trac.wxwidgets.org/ticket/13279
	"""
	if not wx.GetApp():
		dummyApp = wx.App(redirect=False)
		
	sdc = wx.ScreenDC()
# 	width = sdc.GetSize().GetWidth()
# 	height = sdc.GetSize().GetHeight()
# 	bmp = wx.EmptyBitmap(width, height)
# 	mdc = wx.MemoryDC(bmp)
# 	mdc.Blit(0, 0, width, height, sdc, 0, 0)
# 	del mdc
	bmp = sdc.GetAsBitmap()
	
	#load into PIL image
	tempfile = NamedTemporaryFile(prefix='qs_', delete=False)
	tempfile.close()
	bmp.SaveFile(tempfile.name, wx.BITMAP_TYPE_BMP)
	im = Image.open(tempfile.name)
	im.load()
	os.remove(tempfile.name)
	
	return (im, None)

def snap_win(showcursor = True):
    import win32gui,  win32ui,  win32con, win32api, win32clipboard
    hwnd = win32gui.GetDesktopWindow()
    
    # get complete virtual screen including all monitors
    SM_XVIRTUALSCREEN = 76
    SM_YVIRTUALSCREEN = 77
    SM_CXVIRTUALSCREEN = 78
    SM_CYVIRTUALSCREEN = 79
    w = vscreenwidth = win32api.GetSystemMetrics(SM_CXVIRTUALSCREEN)
    h = vscreenheigth = win32api.GetSystemMetrics(SM_CYVIRTUALSCREEN)
    l = vscreenx = win32api.GetSystemMetrics(SM_XVIRTUALSCREEN)
    t = vscreeny = win32api.GetSystemMetrics(SM_YVIRTUALSCREEN)
    r = l + w
    b = t + h
     
    hwndDC = win32gui.GetWindowDC(hwnd)
    mfcDC  = win32ui.CreateDCFromHandle(hwndDC)
    saveDC = mfcDC.CreateCompatibleDC()

    saveBitMap = win32ui.CreateBitmap()
    saveBitMap.CreateCompatibleBitmap(mfcDC, w, h)
    saveDC.SelectObject(saveBitMap)
    saveDC.BitBlt((0, 0), (w, h),  mfcDC,  (l, t),  win32con.SRCCOPY)

    #add cursor
    if showcursor:
        # curFlags, curH, (curX, curY) = win32gui.GetCursorInfo()
        # saveDC.DrawIcon((curX, curY), curH)
        curFlags, curH, (curX, curY) = win32gui.GetCursorInfo()
        curWidth = win32api.GetSystemMetrics(win32con.SM_CXCURSOR)
        curHeight = win32api.GetSystemMetrics(win32con.SM_CYCURSOR)
        relCurX = curX - l - curWidth/4
        relCurY = curY - t - curHeight/4
        saveDC.DrawIcon((relCurX, relCurY), curH)
    
    #load into PIL image
    """http://stackoverflow.com/questions/4199497/image-frombuffer-with-16-bit-image-data"""
    """bmpinfo = saveBitMap.GetInfo()
    bmpstr = saveBitMap.GetBitmapBits(True)
    try:
        stride = 2*bmpinfo['bmWidth']
        #stride = (stride + 3/4) * 4
        im = Image.frombuffer(
            'RGB',
            (bmpinfo['bmWidth'], bmpinfo['bmHeight']),
            bmpstr, 'raw', 'RGBX', 0, 1)
    except ValueError, e: #32 bits didn't work, trying 16 bits
        print e
        im = Image.frombuffer(
            'RGB',
            (bmpinfo['bmWidth'], bmpinfo['bmHeight']),
            bmpstr, 'raw', 'BGR', 0, 1)"""
    
    """Avoid dealing with multiple color depth (the above doesn't work with 24 bits
    and was not tested with <16 bits), write the bitmap to disk then load it from
    there"""
    tempfile = NamedTemporaryFile(prefix='qs_', delete=False)
    tempfile.close()
    saveBitMap.SaveBitmapFile(saveDC, tempfile.name)
    im = Image.open(tempfile.name)
    im.load()
    os.remove(tempfile.name)

    win32gui.DeleteObject(saveBitMap.GetHandle())
    saveDC.DeleteDC()
    mfcDC.DeleteDC()
    win32gui.ReleaseDC(hwnd, hwndDC)

    return (im, None)
    
if __name__ == "__main__":
    (im, bitmap) = snap()
    im.save('screenshot.png', 'PNG')